package com.ash.threadpool;

import java.lang.Thread.UncaughtExceptionHandler;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Observable;
import java.util.Observer;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Throwables;
import com.google.common.util.concurrent.ThreadFactoryBuilder;

/**
 * 负载线程池,分发线程为单线程线程池，保证了顺序。每个任务线程池都为单线程池<br>
 * <li>优先安排任务到包含该任务的key的线程中</li>
 * <li>否则分配任务到任务最少的线程池中</li>
 * 
 * @author Ashin
 *
 */
public class LoadThreadPool implements ILoadThreadPool, LoadThreadPoolMBean {

    private static final Logger logger = LoggerFactory.getLogger(LoadThreadPool.class);

    private String threadName;
    private int threadPoolNum;
    private int queueLenPerPool;

    // 正在处理的任务
    private final BlockingQueue<String> handlingTaskQueue = new LinkedBlockingQueue<String>();
    // key对应的线程池编号
    private Map<String, Integer> keyAndExcutorNoMap = new HashMap<String, Integer>();
    // 编号和线程池的映射
    private Map<Integer, ThreadPoolExecutor> taskExecutorMap;

    private ThreadPoolExecutor relayExcutor;

    public LoadThreadPool() {
        super();
    }

    public LoadThreadPool(String threadName, int threadPoolNum, int queueLenPerPool) {
        this.threadName = threadName;
        this.threadPoolNum = threadPoolNum;
        this.queueLenPerPool = queueLenPerPool;

        relayExcutor = initRelayExcutor();

        taskExecutorMap = initThreadPool();
    }

    private ThreadPoolExecutor initRelayExcutor() {
        ThreadFactoryBuilder tfb = new ThreadFactoryBuilder();
        tfb.setDaemon(true).setNameFormat("lp-relay-%d").setUncaughtExceptionHandler(new UncaughtExceptionHandler() {

            @Override
            public void uncaughtException(Thread t, Throwable e) {
                logger.error("relay exception :{}, task : {} ", Throwables.getStackTraceAsString(e), t);
            }
        });

        ThreadPoolExecutor relayExcutor = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS,
                new LinkedBlockingQueue<Runnable>(), tfb.build());

        return relayExcutor;
    }

    private Map<Integer, ThreadPoolExecutor> initThreadPool() {
        Map<Integer, ThreadPoolExecutor> executorServiceMap = new HashMap<Integer, ThreadPoolExecutor>();
        for (int i = 0; i < threadPoolNum; i++) {
            executorServiceMap.put(i, newSingleExecutor(String.format(this.threadName
                    + "-%d", i), new LinkedBlockingQueue<Runnable>(queueLenPerPool)));
        }

        return executorServiceMap;
    }

    private ThreadPoolExecutor newSingleExecutor(String threadName, BlockingQueue<Runnable> queue) {
        ThreadFactoryBuilder tfb = new ThreadFactoryBuilder();
        tfb.setDaemon(true).setNameFormat(threadName).setUncaughtExceptionHandler(new UncaughtExceptionHandler() {

            @Override
            public void uncaughtException(Thread t, Throwable e) {
                logger.error("load thread executor exception :{}, task : {} ", Throwables.getStackTraceAsString(e), t);
            }
        });

        ThreadPoolExecutor excutor = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, queue, tfb.build());

        return excutor;
    }

    @Override
    public void excuete(final Task task) {
        relayExcutor.execute(new Runnable() {

            @Override
            public void run() {
                logger.debug("receive task {}", task);
                // 注册任务完成监听器
                addTaskListener(task);

                boolean contains = handlingTaskQueue.contains(task.getKey());
                Integer threadNo = 0;
                if (contains) {
                    threadNo = keyAndExcutorNoMap.get(task.getKey());

                } else {
                    // 获取任务最少的线程池中
                    threadNo = fetchMinTaskOfThreadNo(taskExecutorMap);
                    keyAndExcutorNoMap.put(task.getKey(), threadNo);

                    handlingTaskQueue.add(task.getKey());
                }

                ThreadPoolExecutor executorService = taskExecutorMap.get(threadNo);
                executorService.execute(task);

                logger.debug("execute task {}", task);

            }

        });
    }

    private int fetchMinTaskOfThreadNo(Map<Integer, ThreadPoolExecutor> executorServiceMap) {
        int minThreadNo = 0;
        int minTaskQueue = Integer.MAX_VALUE;
        for (Entry<Integer, ThreadPoolExecutor> entry : executorServiceMap.entrySet()) {
            int queueSize = entry.getValue().getQueue().size();
            if (queueSize < minTaskQueue) {
                minTaskQueue = queueSize;
                minThreadNo = entry.getKey();
            }
        }
        return minThreadNo;
    }

    private void addTaskListener(final Task task) {
        task.addObserver(new Observer() {

            @Override
            public void update(Observable o, Object arg) {
                if (o instanceof Task) {
                    handlingTaskQueue.remove(((Task) o).getKey());

                    logger.debug("complete task {}, is handing queue {}", task, handlingTaskQueue);
                }
            }
        });
    }

    @Override
    public void shutdown() {
        for (Entry<Integer, ThreadPoolExecutor> entry : taskExecutorMap.entrySet()) {
            entry.getValue().shutdown();
        }

        if (relayExcutor != null) {
            relayExcutor.shutdown();
        }
    }

    @Override
    public String getAllExecutorState() {
        StringBuilder sb = new StringBuilder();
        for (Entry<Integer, ThreadPoolExecutor> entry : taskExecutorMap.entrySet()) {
            sb.append(entry.getKey()).append(" : ").append(entry.getValue().getQueue().size()).append("; ");
        }
        return sb.toString();
    }

}
